# -*- coding: utf-8 -*-
"""
Rerank模型基础抽象类
定义所有rerank模型的统一接口
"""

from abc import ABC, abstractmethod
from typing import List, Tuple, Dict, Any, Optional
import time
from loguru import logger


class BaseReranker(ABC):
    """
    Rerank模型基础抽象类
    
    所有rerank模型都应该继承此类并实现相应的抽象方法
    """
    
    def __init__(self, model_name: str, **kwargs):
        """
        初始化rerank模型
        
        Args:
            model_name: 模型名称
            **kwargs: 其他配置参数
        """
        self.model_name = model_name
        self.config = kwargs
        self._model = None
        self._is_loaded = False
        
    @abstractmethod
    def load_model(self) -> None:
        """
        加载模型
        
        子类必须实现此方法来加载具体的模型
        """
        pass
    
    @abstractmethod
    def rerank(self, query: str, candidates: List[str], top_k: Optional[int] = None) -> List[Tuple[str, float]]:
        """
        对候选文档进行重排序
        
        Args:
            query: 查询文本
            candidates: 候选文档列表
            top_k: 返回前k个结果，如果为None则返回所有结果
            
        Returns:
            List[Tuple[str, float]]: 重排序后的(文档, 分数)列表，按分数降序排列
        """
        pass
    
    def batch_rerank(self, queries: List[str], candidates_list: List[List[str]], 
                    top_k: Optional[int] = None) -> List[List[Tuple[str, float]]]:
        """
        批量重排序
        
        Args:
            queries: 查询文本列表
            candidates_list: 每个查询对应的候选文档列表
            top_k: 返回前k个结果
            
        Returns:
            List[List[Tuple[str, float]]]: 每个查询的重排序结果
        """
        results = []
        for query, candidates in zip(queries, candidates_list):
            result = self.rerank(query, candidates, top_k)
            results.append(result)
        return results
    
    def evaluate_speed(self, query: str, candidates: List[str], num_runs: int = 10) -> Dict[str, float]:
        """
        评估模型速度性能
        
        Args:
            query: 测试查询
            candidates: 测试候选文档
            num_runs: 运行次数
            
        Returns:
            Dict[str, float]: 包含平均时间、最小时间、最大时间等统计信息
        """
        if not self._is_loaded:
            self.load_model()
            
        times = []
        for _ in range(num_runs):
            start_time = time.time()
            self.rerank(query, candidates)
            end_time = time.time()
            times.append(end_time - start_time)
        
        return {
            'avg_time': sum(times) / len(times),
            'min_time': min(times),
            'max_time': max(times),
            'total_time': sum(times),
            'num_runs': num_runs
        }
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            Dict[str, Any]: 模型信息字典
        """
        return {
            'model_name': self.model_name,
            'is_loaded': self._is_loaded,
            'config': self.config
        }
    
    def __str__(self) -> str:
        return f"{self.__class__.__name__}(model_name='{self.model_name}')"
    
    def __repr__(self) -> str:
        return self.__str__()